from utils.log_helper import log_init
import torch
from model.TranslationModel import TranslationModel

if __name__ == '__main__':
    logger = log_init('test_TranslationModel')
    src_len = 3
    batch_size = 2
    dmodel = 8
    tgt_len = 3
    num_head = 2

    # 转换成 [src_len, batch_size]
    src = torch.tensor([[4, 3, 0],
                        [5, 7, 8]]).transpose(0, 1)
    src_key_padding_mask = torch.tensor([[True, True, False],
                                         [True, True, True]])
    tgt = torch.tensor([[1, 3, 5],
                        [1, 6, 0]]).transpose(0, 1)
    tgt_key_padding_mask = torch.tensor([
        [True, True, True],
        [True, True, False]])
    trans_model = TranslationModel(src_vocab_size=9, tgt_vocab_size=9,
                                   d_model=dmodel, nhead=num_head,
                                   num_encoder_layers=1,
                                   num_decoder_layers=1,
                                   dim_feedforward=30, dropout=0.1)
    tgt_mask = trans_model.my_transformer.generate_square_subsequent_mask(tgt_len)

    # logits shape [tgt_len, batch_size, tgt_vocab_size]
    logits = trans_model(src, tgt=tgt, tgt_mask=tgt_mask,
                         src_key_padding_mask=src_key_padding_mask,
                         tgt_key_padding_mask=tgt_key_padding_mask,
                         memory_key_padding_mask=src_key_padding_mask)
    logger.debug(f'output shape = {logits.shape}')
